package cn.ddiancan.xddcloud.common.loadbalancer;

import java.util.List;
import java.util.Objects;
import java.util.Random;
import java.util.concurrent.atomic.AtomicInteger;

import cn.ddiancan.xddcloud.common.constants.XddHeaderConstant;
import cn.ddiancan.xddcloud.common.context.CloudRequestContext;
import cn.ddiancan.xddcloud.common.context.XddCloudContextUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.DefaultResponse;
import org.springframework.cloud.client.loadbalancer.EmptyResponse;
import org.springframework.cloud.client.loadbalancer.Request;
import org.springframework.cloud.client.loadbalancer.Response;
import org.springframework.cloud.loadbalancer.core.NoopServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.core.SelectedInstanceCallback;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Mono;

public class XddcloudCustomLoadbalancer implements ReactorServiceInstanceLoadBalancer {

    private static final Logger log = LoggerFactory.getLogger(XddcloudCustomLoadbalancer.class);

    final AtomicInteger position;

    final String serviceId;

    ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider;

    public XddcloudCustomLoadbalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider,
        String serviceId) {
        this(serviceInstanceListSupplierProvider, serviceId, new Random().nextInt(1000));
    }

    public XddcloudCustomLoadbalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider,
        String serviceId, int seedPosition) {
        this.serviceId = serviceId;
        this.serviceInstanceListSupplierProvider = serviceInstanceListSupplierProvider;
        this.position = new AtomicInteger(seedPosition);
    }

    @SuppressWarnings("rawtypes")
    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        ServiceInstanceListSupplier supplier =
            serviceInstanceListSupplierProvider.getIfAvailable(NoopServiceInstanceListSupplier::new);
        return supplier.get(request).next()
            .map(serviceInstances -> processInstanceResponse(supplier, serviceInstances));
    }

    private Response<ServiceInstance> processInstanceResponse(ServiceInstanceListSupplier supplier,
        List<ServiceInstance> serviceInstances) {
        Response<ServiceInstance> serviceInstanceResponse = getInstanceResponse(serviceInstances);
        if (supplier instanceof SelectedInstanceCallback && serviceInstanceResponse.hasServer()) {
            ((SelectedInstanceCallback) supplier).selectedServiceInstance(serviceInstanceResponse.getServer());
        }
        return serviceInstanceResponse;
    }

    private Response<ServiceInstance> getInstanceResponse(List<ServiceInstance> instances) {
        log.info("-------------->source service:【{}】,target service:【{}】", XddCloudContextUtils.getApplication(),
            serviceId);
        int instanceSize = instances.size();
        log.info("-------------->xddcloud loadbalancer 【{}】 instances size:【{}】", serviceId, instanceSize);
        if (instances.isEmpty()) {
            log.warn("No servers available for service: {}", serviceId);
            return new EmptyResponse();
        }
        if (instanceSize == 1) {
            return new DefaultResponse(instances.get(0));
        }
        if (CloudRequestContext.isCurrentThreadInitialized()) {
            String first = Objects.requireNonNull(CloudRequestContext.getContextForCurrentThread()).getOutGoing()
                .getFirst(XddHeaderConstant.X_INSTANCE_ROUTE);
            if (StringUtils.hasText(first)) {
                return new DefaultResponse(
                    instances.stream().filter(serviceInstance -> serviceInstance.getInstanceId().equals(first))
                        .findFirst().orElse(instances.get(0)));
            }
        }
        log.info("-------------->xddcloud loadbalancer current position:【{}】", position);
        // 计算加权轮询的权重位置
        int pos = this.position.incrementAndGet() & Integer.MAX_VALUE;
        // 根据权重进行加权轮询
        int totalWeight = instances.stream().mapToInt(instance -> {
            String weightStr = instance.getMetadata().get("weight");
            weightStr = StringUtils.hasText(weightStr) ? weightStr : "1";  // 默认为1
            return Integer.parseInt(weightStr);
        }).sum();
        int currentWeight = pos % totalWeight;  // 当前请求的加权位置
        // 根据权重选择实例
        for (ServiceInstance instance : instances) {
            String weightStr = instance.getMetadata().get("weight");
            weightStr = StringUtils.hasText(weightStr) ? weightStr : "1";
            int weight = Integer.parseInt(weightStr);
            // 判断是否选择当前实例
            if (currentWeight < weight) {
                log.info("-------------->最终选择实例：【{}】,实例信息：【{}】", instance.getInstanceId(),
                    (instance.getHost() + instance.getPort()));
                return new DefaultResponse(instance);
            }
            currentWeight -= weight;  // 减去已选择的权重
        }
        // 如果没有找到合适的实例，默认选择
        ServiceInstance defaultInstance = instances.get(pos % instanceSize);
        log.info("-------------->最终选择实例：【{}】,实例信息：【{}】", defaultInstance.getInstanceId(),
            (defaultInstance.getHost() + defaultInstance.getPort()));
        return new DefaultResponse(defaultInstance);
    }

}
